import os
import tkinter as tk
from tkinter import filedialog, messagebox
from rich.tree import Tree
from rich.console import Console
from tkinter.scrolledtext import ScrolledText


def build_tree(
    directory, parent, max_depth, exclude, current_depth=0, use_symbols=True
):
    """
    递归构建目录树。

    :param directory: 要生成树的目录路径
    :param parent: rich.tree.Tree 的父节点
    :param max_depth: 最大递归深度，-1 表示不限制深度
    :param exclude: 排除的文件或目录列表
    :param current_depth: 当前递归深度
    :param use_symbols: 是否在输出中使用表情符号
    """
    # 如果当前深度超过最大深度且最大深度不为 -1，停止递归
    if max_depth != -1 and current_depth >= max_depth:
        return

    for item in os.listdir(directory):
        # 如果目录或文件在排除列表中，跳过
        if item in exclude:
            continue

        full_path = os.path.join(directory, item)
        if os.path.isdir(full_path):
            # 根据选项添加目录到树中，并递归构建子树
            branch = parent.add(f"📂 {item}" if use_symbols else f"{item}")
            build_tree(
                full_path, branch, max_depth, exclude, current_depth + 1, use_symbols
            )
        else:
            # 根据选项添加文件到树中
            parent.add(f"📄 {item}" if use_symbols else f"{item}")


def generate_tree():
    """
    生成目录树并显示在文本框中。
    """
    # 获取用户输入的目录路径
    directory = directory_var.get()
    if not directory:
        messagebox.showerror("错误", "请选择目录路径！")
        return

    try:
        # 获取用户输入的最大深度
        max_depth = int(depth_var.get())
    except ValueError:
        messagebox.showerror("错误", "最大深度必须是整数！")
        return

    # 获取要排除的文件或目录列表
    exclude = exclude_text.get("1.0", tk.END).strip().split("\n")

    # 检查是否启用表情符号
    use_symbols = symbol_var.get() == 1

    # 创建目录树根节点
    tree = Tree(f"{os.path.basename(directory)}")
    # 构建目录树
    build_tree(directory, tree, max_depth, exclude, use_symbols=use_symbols)

    # 使用 rich.console.Console 显示树
    console = Console(record=True)
    console.print(tree)
    # 清空并插入新的树结构到输出框
    output_text.delete("1.0", tk.END)
    output_text.insert(tk.END, console.export_text())


def select_directory():
    """
    打开目录选择对话框并设置用户选择的目录路径。
    """
    selected_directory = filedialog.askdirectory()
    if selected_directory:
        directory_var.set(selected_directory)


# 创建主窗口
root = tk.Tk()
root.title("目录树生成器")

# Frame for selecting directory
frame_top = tk.Frame(root)
frame_top.pack(fill=tk.X, padx=10, pady=5)

# 创建“选择目录”输入框和按钮
tk.Label(frame_top, text="选择目录:").grid(row=0, column=0, sticky="w")
directory_var = tk.StringVar()
directory_entry = tk.Entry(frame_top, textvariable=directory_var)
directory_entry.grid(row=0, column=1, sticky="we", padx=5)
tk.Button(frame_top, text="浏览", command=select_directory).grid(
    row=0, column=2, sticky="e"
)

# 配置列的权重，确保输入框可以动态调整宽度
frame_top.columnconfigure(1, weight=1)
frame_top.columnconfigure(2, weight=0)

# Frame for max depth
frame_depth = tk.Frame(root)
frame_depth.pack(fill=tk.X, padx=10, pady=5)

# 创建“最大查找深度”输入框
tk.Label(frame_depth, text="最大查找深度:").pack(side=tk.LEFT)
depth_var = tk.StringVar(value="3")
depth_entry = tk.Entry(frame_depth, textvariable=depth_var, width=10)
depth_entry.pack(side=tk.LEFT, padx=5)

# Frame for symbol option
frame_symbols = tk.Frame(root)
frame_symbols.pack(fill=tk.X, padx=10, pady=5)

# 创建是否使用表情符号的复选框
tk.Label(frame_symbols, text="输出符号文字:").pack(side=tk.LEFT)
symbol_var = tk.IntVar(value=1)
symbol_checkbox = tk.Checkbutton(frame_symbols, variable=symbol_var)
symbol_checkbox.pack(side=tk.LEFT, padx=5)

# Frame for exclude list
frame_exclude = tk.Frame(root)
frame_exclude.pack(fill=tk.X, padx=10, pady=5)

# 创建“排除列表”输入框
tk.Label(frame_exclude, text="要排除的目录或文件，以换行分隔:").pack(anchor=tk.W)
exclude_text = ScrolledText(frame_exclude, height=5)
exclude_text.pack(fill=tk.X, padx=5, pady=5)

# Frame for output
frame_output = tk.Frame(root)
frame_output.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)

# 创建“目录树”输出框
tk.Label(frame_output, text="目录树:").pack(anchor=tk.W)
output_text = ScrolledText(frame_output, height=15)
output_text.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

# 生成目录树按钮
generate_button = tk.Button(root, text="生成目录树", command=generate_tree)
generate_button.pack(pady=10)
# 获取屏幕的宽度和高度
screen_width = root.winfo_screenwidth()
screen_height = root.winfo_screenheight()

# 获取窗口的初始宽度和高度
window_width = 800  # 设置一个初始的宽度
window_height = 600  # 设置一个初始的高度

# 计算居中的位置
position_top = int((screen_height - window_height) / 2)
position_left = int((screen_width - window_width) / 2)

# 设置窗口位置并允许调整大小
root.geometry(f"{window_width}x{window_height}+{position_left}+{position_top}")
root.resizable(True, True)  # 允许窗口大小变化

# 运行主事件循环
root.mainloop()
